import flask
from flask import request
import urllib
from PIL import Image
import io
import numpy as np
import matplotlib.pyplot as plt
import data
import torch
import torch.nn as nn


class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(1, 6, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(6, 16, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        self.model2 = nn.Sequential(
            nn.Linear(16 * 5 * 5, 120),
            nn.ReLU(),
            nn.Linear(120, 84),
            nn.ReLU(),
            nn.Linear(84, 3980)
        )

    def forward(self, wa):
        out = self.model(wa)
        return self.model2(out.view(out.shape[0], -1))


app = flask.Flask(__name__)
net = torch.load('./model2').cpu()


@app.route('/catch', methods=['POST'])
def catch():
    response = urllib.request.urlopen(request.form['id'])
    buf = io.BytesIO()
    buf.write(response.file.read())
    arr = np.asarray(Image.open(buf))
    height, width = arr.shape[0], arr.shape[1]
    new_arr = np.zeros((height, width))
    for i in range(height):
        for j in range(width):
            new_arr[i][j] = 255 - arr[i][j][3]
    new_data = data.convert_from_numpy(new_arr, 32, 32)
    in_data = torch.Tensor(new_data).view(1, 1, 32, 32).float()
    with torch.no_grad():
        y_hat = net(in_data)
    return data.get_ans(y_hat.argmax().item())


@app.route('/')
def index():
    return flask.render_template('index.html')


if __name__ == '__main__':
    app.run()
